#!/usr/bin/env python

import rospy
import tf2_ros
import geometry_msgs.msg
import std_msgs.msg
import tf.transformations as tft

def main():
    rospy.init_node('publish_body2d_node')

    # Create TF buffer and listener
    tf_buffer = tf2_ros.Buffer()
    tf_listener = tf2_ros.TransformListener(tf_buffer)

    # Create broadcaster
    tf_broadcaster = tf2_ros.TransformBroadcaster()

    rate = rospy.Rate(10.0)  # Publish at 10 Hz

    while not rospy.is_shutdown():
        try:
            # Lookup transform from map to body
            trans = tf_buffer.lookup_transform(
                "map",
                "body",
                rospy.Time(0),
                rospy.Duration(1.0)
            )

            # Extract translation (keep as is)
            translation = [trans.transform.translation.x,
                           trans.transform.translation.y,
                           trans.transform.translation.z]

            # Extract quaternion and convert to Euler angles
            quat = [trans.transform.rotation.x,
                    trans.transform.rotation.y,
                    trans.transform.rotation.z,
                    trans.transform.rotation.w]
            roll, pitch, yaw = tft.euler_from_quaternion(quat)

            # Remove pitch and roll, keep yaw
            new_roll = 0.0
            new_pitch = 0.0
            new_yaw = yaw

            # Convert back to quaternion
            new_quat = tft.quaternion_from_euler(new_roll, new_pitch, new_yaw)

            # Create new transform message for map to body_2d
            t = geometry_msgs.msg.TransformStamped(
                header=std_msgs.msg.Header(
                    frame_id="map",
                    stamp=rospy.Time.now()
                ),
                child_frame_id="body_2d",
                transform=geometry_msgs.msg.Transform(
                    translation=geometry_msgs.msg.Vector3(
                        x=translation[0],
                        y=translation[1],
                        z=0.0
                    ),
                    rotation=geometry_msgs.msg.Quaternion(
                        x=new_quat[0],
                        y=new_quat[1],
                        z=new_quat[2],
                        w=new_quat[3]
                    )
                )
            )

            # Broadcast the transform
            tf_broadcaster.sendTransform(t)

        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            rospy.logwarn("TF lookup failed: %s" % e)

        rate.sleep()

if __name__ == '__main__':
    main()